## Initialize
from ochumanApi.ochuman import OCHuman
# <Filter>:
#      None(default): load all. each has a bbox. some instances have keypoint and some have mask annotations.
#            images: 5081, instances: 13360
#     'kpt&segm' or 'segm&kpt': only load instances contained both keypoint and mask annotations (and bbox)
#            images: 4731, instances: 8110
#     'kpt|segm' or 'segm|kpt': load instances contained either keypoint or mask annotations (and bbox)
#            images: 5081, instances: 10375
#     'kpt' or 'segm': load instances contained particular kind of annotations (and bbox)
#            images: 5081/4731, instances: 10375/8110
# ochuman = OCHuman(AnnoFile='/root/data/OCHuman/ochuman.json', Filter='kpt&segm')
# ochuman = OCHuman(AnnoFile='/root/data/OCHuman/ochuman.json', Filter='segm')
ochuman = OCHuman(AnnoFile='/root/data/OCHuman/ochuman.json', Filter=None)
image_ids = ochuman.getImgIds()
print ('Total images: %d'%len(image_ids))

## Load annotations & Visualize
import cv2, os
import matplotlib
import matplotlib.pyplot as plt
# %matplotlib inline
plt.rcParams['figure.figsize'] = (15, 15)
import ochumanApi.vis as vistool
from ochumanApi.ochuman import Poly2Mask

colors = [[255, 0, 0],
         [255, 255, 0],
         [0, 255, 0],
         [0, 255, 255],
         [0, 0, 255],
         [255, 0, 255]]

ImgDir = '/root/data/OCHuman/images/'

for image_id in image_ids:
    imgs = ochuman.loadImgs(imgIds=[image_id])
    data = imgs[0]

    img = cv2.imread(os.path.join(ImgDir, data['file_name']))
    height, width = data['height'], data['width']

    for i, anno in enumerate(data['annotations']):
        bbox = anno['bbox']
        kpt = anno['keypoints']
        segm = anno['segms']
        max_iou = anno['max_iou']

        img = vistool.draw_bbox(img, bbox, thickness=3, color=colors[i%len(colors)])
        if segm is not None:
            mask = Poly2Mask(segm)
            img = vistool.draw_mask(img, mask, thickness=3, color=colors[i%len(colors)])
        if kpt is not None:
            img = vistool.draw_skeleton(img, kpt, connection=None, colors=colors[i%len(colors)], bbox=bbox)

    cv2.imshow('image', img)
    key = cv2.waitKey()
    if key == 27 or key == ord('q') or key == ord('Q'):
        break

# imgs = ochuman.loadImgs(imgIds=[image_ids[2]])
# data = imgs[0]
#
# img = cv2.imread(os.path.join(ImgDir, data['file_name']))
# height, width = data['height'], data['width']
#
# for i, anno in enumerate(data['annotations']):
#     bbox = anno['bbox']
#     kpt = anno['keypoints']
#     segm = anno['segms']
#     max_iou = anno['max_iou']
#
#     img = vistool.draw_bbox(img, bbox, thickness=3, color=colors[i%len(colors)])
#     if segm is not None:
#         mask = Poly2Mask(segm)
#         img = vistool.draw_mask(img, mask, thickness=3, color=colors[i%len(colors)])
#     if kpt is not None:
#         img = vistool.draw_skeleton(img, kpt, connection=None, colors=colors[i%len(colors)], bbox=bbox)
#
# plt.imshow(img[:,:,::-1])
# plt.axis('off')
# plt.show()

## Convert to coco format.
# <maxIouRange>:
#     (0.0, 1.0) means all instances
#     (0.5, 0.75) means Moderate instances
#     (0.75, 1.0) means Hard instances
ochuman.toCocoFormart(subset='val', maxIouRange=(0.0, 1.0), save_dir='./')
ochuman.toCocoFormart(subset='test', maxIouRange=(0.0, 1.0), save_dir='./')
